
/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: GodK
 * Date: 2024-05-07
 * Time: 10:55
 */
public class RBTree {
    static class RBTreeNode{
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public COLOR color;
        public RBTreeNode(int val){
            this.val = val;
            this.color = COLOR.RED;
        }
    }
    public RBTreeNode root;

    public boolean insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if(root == null){
            root = node;
            root.color = COLOR.BLACK;
            return true;
        }
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while(cur != null){
            if(val > cur.val){
                parent = cur;
                cur = cur.right;
            }else if(val < cur.val){
                parent = cur;
                cur = cur.left;
            }else {
                return false;
            }
        }
        if(parent.val > val){
            parent.left = node;
        }else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        while(parent != null && parent.color == COLOR.RED){
            RBTreeNode grandFather = parent.parent;
            if(grandFather.left == parent){
                RBTreeNode uncle = grandFather.right;
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    cur = grandFather;
                    parent = cur.parent;
                }else {
                    if(parent.right == cur){
                        rotateLeft(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }
                    rotateRight(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }else {
                RBTreeNode uncle = grandFather.left;
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    cur = grandFather;
                    parent = cur.parent;
                }else {
                    if(parent.left == cur){
                        rotateRight(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }
                    rotateLeft(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }
        }
        root.color = COLOR.BLACK;
        cur = root;
        int blackNum = 0;
        while(cur != null){
            if(cur.color == COLOR.BLACK){
                blackNum++;
            }
            cur = cur.left;
        }
        return isRBTree() && checkRedColor(root) && checkBlackNum(root,0,blackNum);
    }

    //左单旋
    private void rotateLeft(RBTreeNode parent) {
        RBTreeNode SubR = parent.right;
        RBTreeNode SubRL = SubR.left;
        parent.right = SubRL;
        SubR.left = parent;
        if (SubRL != null) {
            SubRL.parent = parent;
        }
        RBTreeNode pParent = parent.parent;
        parent.parent = SubR;
        if(root == parent) {
            root = SubR;
            root.parent = null;
        }else{
            if (pParent.left == parent) {
                pParent.left = SubR;
            }else {
                pParent.right = SubR;
            }
            SubR.parent = pParent;
        }
    }
    //右单旋
    private void rotateRight(RBTreeNode parent){
        RBTreeNode SubL = parent.left;
        RBTreeNode SubLR = SubL.right;
        parent.left = SubLR;
        SubL.right = parent;
        if(SubLR != null){
            SubLR.parent = parent;
        }
        RBTreeNode pParent = parent.parent;
        parent.parent = SubL;
        if(root == parent){
            root = SubL;
            root.parent = null;
        }else {
            if(pParent.left == parent){
                pParent.left = SubL;
            }else {
                pParent.right = SubL;
            }
            SubL.parent = pParent;
        }
    }

    public boolean isRBTree(){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.RED){
            System.out.println("根节点不是黑色");
        }
        return true;
    }

    private boolean checkRedColor(RBTreeNode root){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.RED){
            RBTreeNode parent = root.parent;
            if (parent.color == COLOR.RED) {
                System.out.println("连续出现两个红色节点");
                return false;
            }
        }
        return checkRedColor(root.left) && checkRedColor(root.right);
    }

    private boolean checkBlackNum(RBTreeNode root,int pathBlackNum,int blackNum){
        if(root ==null){
            return true;
        }
        if(root.color == COLOR.BLACK){
            pathBlackNum++;
        }
        if(root.left == null && root.right == null){
            if(pathBlackNum != blackNum){
                System.out.println("每条路径上的黑色节点数不一样");
                return false;
            }
        }
        return checkBlackNum(root.left,pathBlackNum,blackNum) && checkBlackNum(root.right,pathBlackNum,blackNum);
    }
}
